import numpy as np
import matplotlib.pyplot as plt
import os
import sys
from check import *
from utils2 import *
from utils3 import *
import shutil

# dir = r"D:\data_cmp\airfoil_csv"

def loop_dataset_prt(f):
    count = 0  # 总数
    read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip"
    save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_clip"

    # leading count
    count0 = 0  # x为0的坐标的个数不为1的个数
    count1 = 0  # 前缘x坐标不为0，需要归一化
    count2 = 0  # 前缘y坐标不为0，需要平移

    # leading: y value
    leading_y = []

    # trailing
    count3 = 0  # 后缘不为2个点的个数
    count4 = 0  # 后缘x坐标不为1的个数，需要归一化
    count5 = 0  # 后缘y坐标不为0的个数，需要旋转

    for file in os.listdir(read_dir):
        file_path = os.path.join(read_dir, file)
        leading, trailing = f(file_path)
        # print(leading, trailing)
        count0 += (False == leading[0])
        count1 += (False == leading[1].any())
        count2 += (False == leading[2].any())
        count3 += (False == trailing[0])
        count4 += (False == trailing[1].any())
        count5 += (False == trailing[2].any())
        count += 1
        if not leading[1].any():
            print("x is not 0", leading[3], leading[4], file)

        if not leading[2].any():
            print("y is not 0", leading[3], leading[4], file)
            leading_y += list(abs(leading[4]))

        if not trailing[0]:
            print("shape is not 2:", trailing[3], trailing[4], file)
            save_path = os.path.join(save_dir, file)

        if not trailing[1].any():
            print("x is not 1:", trailing[3], trailing[4], file)
            save_path = os.path.join(save_dir, file)

        if not trailing[2].any():
            print("y is not 0:", trailing[3], trailing[4], file)

    leading_y = np.asarray(leading_y).flat
    # print("leading_y:", np.min(leading_y), np.max(leading_y), np.mean(leading_y))

    print("total count:", count)
    print(count0)
    print(count1)
    print(count2)
    print(count3)
    print(count4)
    print(count5)

    # plt.bar(np.arange(0, len(leading_y)), leading_y, 10)
    # plt.show()


def loop_dataset_seq():
    '''
        遍历整个数据集，找出所有不按逆时针排列的文件
    :return: 返回不按逆时针排列的文件
    '''
    not_in_order = []
    read_dir = r"D:\data_cmp\airfoil_csv"

    for file in os.listdir(read_dir):
        file_path = os.path.join(read_dir, file)
        res = check_seq(file_path, True)
        if res is False:
            print(res, file_path)
            not_in_order.append(file)

    return not_in_order


def loop_dataset_transform():
    '''
        将翼型数据转化为逆时针序列；并保存在D:\data_cmp\airfoil_csv_op\airfoil_csv_op_1_seq_tmp
        有bug，未修复
    :return:
    '''
    read_dir = r"D:\data_cmp\airfoil_csv"
    save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_1_seq_tmp"
    file_set = loop_dataset_seq()
    for file in file_set:
        save_path = os.path.join(save_dir, file)
        print("save_path:", save_path)

        ordinates = pd.read_csv(os.path.join(read_dir, file)).to_numpy()
        tran_ordinates = transform(ordinates)
        df = pd.DataFrame(tran_ordinates)
        df.columns = ["x", "y"]
        df.to_csv(save_path, index=False)


def loop_dataset_clip1():
    '''
        处理机翼前缘及机翼后缘
    :return:
    '''
    # read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_2_clip_tmp"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_2_clip"
    # read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_3_clip_tmp"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\special"
    read_dir = r"D:\data_cmp\airfoil_csv_op\special"
    save_dir = r"D:\data_cmp\airfoil_csv_op\special"
    # for file in os.listdir(read_dir):
    for file in ["sc1095r8.csv"]:
        file_path = os.path.join(read_dir, file)
        save_path = os.path.join(save_dir, file)
        # print("save_path:", save_path)

        ordinates = pd.read_csv(file_path).to_numpy()

        leading, trailing = check_leading_trailing(file_path)
        _, _, _, x_leading, y_leading, id_leading = leading
        _, _, _, x_trailing, y_trailing, id_trailing = trailing

        # print(file, ": ", end="")
        # print(x_leading, y_leading, id_leading, x_trailing, y_trailing, id_trailing)

        # 前缘
        # 左移x
        if x_leading.shape[0] == 2:
            # print(file, ": ", end="")
            # print(x_leading, y_leading, id_leading, x_trailing, y_trailing, id_trailing)
            if x_leading.all() == 0:
                y_max = np.argmax(y_leading)
                y_min = np.argmin(y_leading)
                ordinates = np.concatenate([ordinates[0:id_leading[0]],
                                            np.expand_dims(ordinates[id_leading[y_min]], axis=0),
                                            ordinates[(id_leading[1] + 1):]], axis=0)
                # print(ordinates1[id_leading[0]: id_leading[0]+2])
            else:
                ordinates = move_x(ordinates, -x_leading[0])
        else:  # x_leading.shape[0] == 1
            ordinates = move_x(ordinates, -x_leading[0])

        # 上/下移y
        if id_leading.shape[0] == 1 and ordinates[id_leading[0]][1] != 0:
            # 整体移
            ordinates = move_y(ordinates, -ordinates[id_leading[0]][1])
        elif id_leading.shape[0] == 2 and ordinates[id_leading][1].any() != 0:
            ordinates[0:id_leading[0] + 1] = move_y(ordinates[0:id_leading[0] + 1], -ordinates[id_leading[0]][1])
            ordinates[id_leading[1]::] = move_y(ordinates[id_leading[1]::], -ordinates[id_leading[1]][1])

        # 去掉重复的前缘顶点
        def fmin(x):
            return np.argwhere(x == np.min(x)).flatten()

        id_leading = fmin(ordinates[:, 0])
        if id_leading.shape[0] == 2:
            ordinates = np.delete(ordinates, id_leading[0], axis=0)

        # 保存
        df = pd.DataFrame(ordinates)
        df.columns = ["x", "y"]
        df.to_csv(save_path, index=False)


def loop_dataset_clip2_v1():
    '''
        处理机翼后缘
    :return:
    '''
    # read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_1_seq_all"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_predata1"
    read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_3_clip_tmp"
    save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_3_clip"

    # read_dir = r"D:\data_cmp\airfoil_csv_op\special"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\special"
    # for file in os.listdir(read_dir):
    for file in ["sc1095r8.csv"]:
        # for file in os.listdir(read_dir):
        file_path = os.path.join(read_dir, file)
        save_path = os.path.join(save_dir, file)
        # print("save_path:", save_path)

        ordinates = pd.read_csv(file_path).to_numpy()

        leading, trailing = check_leading_trailing(file_path)
        _, _, _, x_leading, y_leading, id_leading = leading
        _, _, _, x_trailing, y_trailing, id_trailing = trailing

        # print(file, ": ", end="")
        # print(x_leading, y_leading, id_leading, x_trailing, y_trailing, id_trailing)
        assert id_leading.shape[0] == 1

        # 后缘
        # test
        def fmax(x):
            return np.argwhere(x == np.max(x)).flatten()

        # 更准确定位trailing的位置
        id_trailing1 = fmax(ordinates[0:id_leading[0], 0])
        id_trailing2 = fmax(ordinates[id_leading[0]::, 0]) + id_leading[0]

        # 旋转
        ordinates[0:id_leading[0]] = rotate(ordinates[0:id_leading[0]], ordinates[id_trailing1[0]])
        ordinates[id_leading[0]::] = rotate(ordinates[id_leading[0]::], ordinates[id_trailing2[0]])

        # 归一化(横坐标)
        # 上机翼
        x = ordinates[0:id_leading[0] + 1, 0]
        x_min = np.min(x)
        x_max = np.max(x)
        x = (x - x_min) / (x_max - x_min)
        ordinates[0:id_leading[0] + 1, 0] = x  # 要覆盖到0

        # 下机翼
        x = ordinates[id_leading[0]::, 0]
        x_min = np.min(x)
        x_max = np.max(x)
        x = (x - x_min) / (x_max - x_min)
        ordinates[id_leading[0]::, 0] = x

        # 保存
        df = pd.DataFrame(ordinates)
        df.columns = ["x", "y"]
        df.to_csv(save_path, index=False)


def loop_dataset_clip2_v2():
    '''
        处理机翼后缘（第二版）；这种思路行不通
    :return:
    '''
    read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_3_clip_tmp"
    save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_3_clip_v2"  ###

    for file in os.listdir(read_dir):
        file_path = os.path.join(read_dir, file)
        save_path = os.path.join(save_dir, file)
        print("save_path:", save_path)

        ordinates = pd.read_csv(file_path).to_numpy()

        leading, trailing = check_leading_trailing(file_path)
        _, _, _, x_leading, y_leading, id_leading = leading
        _, _, _, x_trailing, y_trailing, id_trailing = trailing

        # print(file, ": ", end="")
        # print(x_leading, y_leading, id_leading, x_trailing, y_trailing, id_trailing)
        assert id_leading.shape[0] == 1

        # 后缘
        # test
        def fmax(x):
            return np.argwhere(x == np.max(x)).flatten()

        # 更准确定位trailing的位置
        id_trailing1 = fmax(ordinates[0:id_leading[0], 0])
        id_trailing2 = fmax(ordinates[id_leading[0]::, 0]) + id_leading[0]

        # 旋转
        basic1 = ordinates[id_trailing1[0]]
        basic2 = ordinates[id_trailing2[0]]
        mid = (basic1[1] + basic2[1]) / 2
        ordinates[0:id_leading[0]] = rotate(ordinates[0:id_leading[0]],
                                            basic1,
                                            angle=np.arcsin(mid / np.sqrt(basic1[0] ** 2 + basic1[1] ** 2)))

        ordinates[id_leading[0]::] = rotate(ordinates[id_leading[0]::],
                                            basic2,
                                            angle=np.arcsin(mid / np.sqrt(basic2[0] ** 2 + basic2[1] ** 2)))

        # break

        # 归一化(横坐标)
        # 上机翼
        x = ordinates[0:id_leading[0] + 1, 0]
        x_min = np.min(x)
        x_max = np.max(x)
        x = (x - x_min) / (x_max - x_min)
        ordinates[0:id_leading[0] + 1, 0] = x  # 要覆盖到0

        # 下机翼
        x = ordinates[id_leading[0]::, 0]
        x_min = np.min(x)
        x_max = np.max(x)
        x = (x - x_min) / (x_max - x_min)
        ordinates[id_leading[0]::, 0] = x

        # 可视化

        # print(ordinates[id_trailing1[0]], ordinates[id_trailing2[0]])
        # plt.title(file+str(basic1[1])+"_"+str(basic2[1]))
        # plt.scatter(ordinates[:, 0], ordinates[:, 1])
        # plt.show()

        # 保存
        # df = pd.DataFrame(ordinates)
        # df.columns = ["x", "y"]
        # df.to_csv(save_path, index=False)


def loop_dataset_set0(read_dir=r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos",
                      save_dir=r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos"):
    for file in os.listdir(read_dir):
        file_path = os.path.join(read_dir, file)
        save_path = os.path.join(save_dir, file)

        ordinates = pd.read_csv(file_path).to_numpy()
        ordinates = set0(ordinates)

        # 保存
        df = pd.DataFrame(ordinates)
        df.columns = ["x", "y"]
        df.to_csv(save_path, index=False)


def loop_dataset_getlen():
    read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip"
    for file in os.listdir(read_dir):
        file_path = os.path.join(read_dir, file)

        ordinates = pd.read_csv(file_path).to_numpy()
        len = ordinates.shape[0]
        print(file, len)


def loop_dataset_interpolation():
    read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip"
    save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi"

    x = produce_x_v1()
    for file in os.listdir(read_dir):
        read_file_path = os.path.join(read_dir, file)
        save_file_path = os.path.join(save_dir, file)
        print(read_file_path)

        ordinates = pd.read_csv(read_file_path).to_numpy()
        ordinates_interpolation = Cubic_Spline_Interpolation_v1(ordinates, x)
        ordinates_interpolation = pd.DataFrame(ordinates_interpolation)
        ordinates_interpolation.columns = ['x', 'y']
        ordinates_interpolation.to_csv(save_file_path, index=False)


def loop_dataset_interpolation_v2(read_dir=r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip",
                                  save_dir=r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos"):
    '''
        对翼型做三次样条插值（取65+65-1=127个点）
    :return:
    '''
    # read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos"
    # save_pic_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos_show"

    n = 151
    x = produce_x_v8(n)
    for file in os.listdir(read_dir):
    # for file in ['hq1512.csv']:
        read_file_path = os.path.join(read_dir, file)
        save_file_path = os.path.join(save_dir, file)
        print(read_file_path)

        error_count = 0
        ordinates = pd.read_csv(read_file_path).to_numpy()

        ordinates_interpolation = Cubic_Spline_Interpolation_v1(ordinates, x, n)

        ordinates_interpolation = pd.DataFrame(ordinates_interpolation)
        ordinates_interpolation.columns = ['x', 'y']
        ordinates_interpolation.to_csv(save_file_path, index=False)


        # plt.title(file)
        # plt.figure(figsize=(20, 10))
        # plt.plot(ordinates_interpolation[:, 0], ordinates_interpolation[:, 1])
        # plt.scatter(ordinates[:, 0], ordinates[:, 1], c='r', s=20)
        # plt.scatter(ordinates_interpolation[:, 0], ordinates_interpolation[:, 1], c='b', s=5)
        # plt.xlim((0, 1.0))  # 设置横坐标范围
        # plt.ylim((-0.3, 0.3))  # 设置纵坐标范围
        # plt.xticks(np.arange(0, 1, 0.1))
        # plt.gca().set_aspect(1)
        # plt.grid()  # 设置网格
        # plt.show()
        #
        # plt.savefig(os.path.join(save_pic_dir, file[:-3] + "jpg"))
        # print(os.path.join(save_pic_dir, file[:-3] + "jpg"))
        # plt.close()
        #


def loop_dataset_interpolation_v3(read_dir=r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip",
                                  save_dir=r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos"):
    '''
        对翼型做三次样条插值（取65+65-1=127个点）（改进方法）（弧长插值）
    :return:
    '''
    # read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos"
    # save_pic_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos_show"

    for file in os.listdir(read_dir):  # ["D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip\a18.csv"]
        # for file in [r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip\a18sm.csv"]:
        read_file_path = os.path.join(read_dir, file)
        save_file_path = os.path.join(save_dir, file)
        print(read_file_path)

        x = produce_x_v3()
        ordinates = pd.read_csv(read_file_path).to_numpy()
        ordinates_interpolation = Cubic_Spline_Interpolation_v3(ordinates, x)

        # save
        ordinates_interpolation = pd.DataFrame(ordinates_interpolation)
        ordinates_interpolation.columns = ['x', 'y']
        ordinates_interpolation.to_csv(save_file_path, index=False)

        # plt.title(file)
        # plt.figure(figsize=(20, 10))
        # plt.plot(ordinates_interpolation[:, 0], ordinates_interpolation[:, 1])
        # plt.scatter(ordinates[:, 0], ordinates[:, 1], c='r', s=20)
        # plt.scatter(ordinates_interpolation[:, 0], ordinates_interpolation[:, 1], c='b', s=5)
        # plt.xlim((0, 1.0))  # 设置横坐标范围
        # plt.ylim((-0.3, 0.3))  # 设置纵坐标范围
        # plt.xticks(np.arange(0, 1, 0.1))
        # plt.gca().set_aspect(1)
        # plt.grid()  # 设置网格
        # plt.show()
        #
        # plt.close()


def loop_dataset_interpolation_v4(read_dir=r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip",
                                  save_dir=r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos"):
    '''
        对翼型做B样条插值（取65+65-1=129个点）
    :return:
    '''
    # read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos"
    # save_pic_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_5_spi_cos_show"

    error_count = 0
    # read_dir_tmp = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_11_rmleading"
    for file in os.listdir(read_dir):
    # for file in ['ag04.csv']:
        read_file_path = os.path.join(read_dir, file)
        save_file_path = os.path.join(save_dir, file)
        # print(read_file_path)

        x_bspine = produce_x_v6(20001)  # 一定要是奇数
        ordinates = pd.read_csv(read_file_path).to_numpy()
        # B-Spline
        try:
            ordinates_interpolation = B_Spline_Interpolation_v3(ordinates, x_bspine, 151)
                # save
            ordinates_interpolation = pd.DataFrame(ordinates_interpolation)
            ordinates_interpolation.columns = ['x', 'y']
            ordinates_interpolation.to_csv(save_file_path, index=False)
        except:
            error_count += 1
            print("{}: error".format(error_count), file)
            shutil.copy(read_file_path,
                        os.path.join(r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_11_bspline_error2", file))


        # plt.title(file)
        # plt.figure(figsize=(20, 10))
        # plt.plot(ordinates_interpolation[:, 0], ordinates_interpolation[:, 1])
        # plt.scatter(ordinates[:, 0], ordinates[:, 1], c='r', s=20)
        # plt.scatter(ordinates_interpolation[:, 0], ordinates_interpolation[:, 1], c='b', s=5)
        # plt.xlim((0, 1.0))  # 设置横坐标范围
        # plt.ylim((-0.3, 0.3))  # 设置纵坐标范围
        # plt.xticks(np.arange(0, 1, 0.1))
        # plt.gca().set_aspect(1)
        # plt.grid()  # 设置网格
        # plt.show()
        # plt.close()


if __name__ == "__main__":
    # loop_dataset_transform()
    # loop_dataset_prt(check_leading_trailing_v2)
    # loop_dataset_seq()
    # loop_dataset_clip1()
    # loop_dataset_clip2_v1()
    # loop_dataset_set0()
    # loop_dataset_getlen()

    # read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_4_clip1"
    # read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_9_bspline_error1"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_9_bspline"
    # read_dir = r"D:\data_cmp\airfoil_csv_op\NACA4"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\NACA4_bspline"
    read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_11_cbi_org"
    save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_11_cbi"

    # B样条插值
    # loop_dataset_interpolation_v4(read_dir=read_dir, save_dir=save_dir)
    # loop_dataset_set0(read_dir=save_dir, save_dir=save_dir)


    # read_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_9_bspline_error2"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\airfoil_csv_op_9_ssp"
    # read_dir = r"D:\data_cmp\airfoil_csv_op\NACA4"
    # save_dir = r"D:\data_cmp\airfoil_csv_op\NACA4_ssp_129"

    # 三次样条插值
    loop_dataset_interpolation_v2(read_dir=read_dir, save_dir=save_dir)
    loop_dataset_set0(read_dir=save_dir, save_dir=save_dir)

